`include "define.sv"

module matrix_router (
    input clk,
    input reset,
    input [1:0] matrix_type,
    input [`DATA_WIDTH-1:0] a_in [0:31],
    input [`DATA_WIDTH-1:0] b_in [0:31],
    input [`DATA_WIDTH-1:0] pe_a_out [0:31][0:15],
    input [`DATA_WIDTH-1:0] pe_b_out [0:15][0:31],
    input [`DATA_WIDTH-1:0] c_in [0:31][0:31],
    input pe_mult_en_out [0:15][0:15],
    output reg [`DATA_WIDTH-1:0] a [0:31][0:31],
    output reg [`DATA_WIDTH-1:0] b [0:31][0:31],
    //output reg [`DATA_WIDTH-1:0] c [0:31][0:31],
    //output reg c_add_en [0:31][0:31],
    input [31:0] a_fifo_empty,
    input [31:0] b_fifo_empty,
    //input c_fifo_empty [0:31][0:31],
    output reg mult_en [0:15][0:15]
);

    // 声明所有genvar变量
    genvar i_mult, j_mult;  // 用于mult_en生成
    genvar i_a, j_a;        // 用于A数据路由
    genvar i_b, j_b;        // 用于B数据路由
    genvar i_c, j_c;        // 用于C数据路由

    /* // 1. 时序逻辑生成mult_en (使用generate)
    generate
        for (i_mult = 0; i_mult < 16; i_mult = i_mult + 1) begin : gen_mult_en_row
            for (j_mult = 0; j_mult < 16; j_mult = j_mult + 1) begin : gen_mult_en_col
                always @(posedge clk or posedge reset) begin
                    if (reset) begin
                        mult_en[i_mult][j_mult] <= 0;
                    end else begin
                        if (j_mult == 0) begin
                            if (matrix_type == `m8n32k16 && i_mult >= 8) begin
                                mult_en[i_mult][j_mult] <= pe_mult_en_out[i_mult - 8][15];
                            end else begin
                                mult_en[i_mult][j_mult] <= ~a_fifo_empty[i_mult];
                            end
                        end else if (j_mult == 8) begin
                            if (matrix_type == `m32n8k16) begin
                                mult_en[i_mult][j_mult] <= ~a_fifo_empty[i_mult + 16];
                            end else begin
                                mult_en[i_mult][j_mult] <= pe_mult_en_out[i_mult][j_mult - 1];
                            end
                        end else begin
                            mult_en[i_mult][j_mult] <= pe_mult_en_out[i_mult][j_mult - 1];
                        end
                    end
                end
            end
        end
    endgenerate */
generate
    for (i_mult = 0; i_mult < 16; i_mult = i_mult + 1) begin : gen_mult_en_row
        for (j_mult = 0; j_mult < 16; j_mult = j_mult + 1) begin : gen_mult_en_col
            always @(*) begin
                if (reset) begin
                    mult_en[i_mult][j_mult] = 0;
                end else begin
                    if (j_mult == 0) begin
                        if (matrix_type == `m8n32k16 && i_mult >= 8) begin
                            mult_en[i_mult][j_mult] = pe_mult_en_out[i_mult - 8][15];
                        end else begin
                            mult_en[i_mult][j_mult] = ~a_fifo_empty[i_mult];
                        end
                    end else if (j_mult == 8) begin
                        if (matrix_type == `m32n8k16) begin
                            mult_en[i_mult][j_mult] = ~a_fifo_empty[i_mult + 16];
                        end else begin
                            mult_en[i_mult][j_mult] = pe_mult_en_out[i_mult][j_mult - 1];
                        end
                    end else begin
                        mult_en[i_mult][j_mult] = pe_mult_en_out[i_mult][j_mult - 1];
                    end
                end
            end
        end
    end
endgenerate
    // 2. 组合逻辑初始化
    always @(*) begin
        // 初始化所有信号
        for (int i = 0; i < 32; i = i + 1) begin
            for (int j = 0; j < 32; j = j + 1) begin
                a[i][j] = 0;
                b[i][j] = 0;
                //c[i][j] = 0;
                //c_add_en[i][j] = 0;
            end
        end
    end

    // A数据路由逻辑
    generate
        for (i_a = 0; i_a < 16; i_a = i_a + 1) begin : a_data_row
            for (j_a = 0; j_a < 16; j_a = j_a + 1) begin : a_data_col
                always @(*) begin
                    // 处理A数据
                    if (j_a == 0) begin
                        if (matrix_type == `m8n32k16) begin
                            if (i_a < 8) begin
                                a[i_a][j_a] = a_in[i_a];
                            end else if (i_a < 16) begin
                                a[i_a][j_a] = pe_a_out[i_a - 8][15];
                            end
                        end else begin
                            if (i_a < 16) begin
                                a[i_a][j_a] = a_in[i_a];
                            end
                        end
                    end else if (j_a == 8) begin
                        if (matrix_type == `m32n8k16) begin
                            a[i_a][j_a] = a_in[i_a + 16];
                        end else begin
                            a[i_a][j_a] = pe_a_out[i_a][j_a - 1];
                        end
                    end else begin
                        a[i_a][j_a] = pe_a_out[i_a][j_a - 1];
                    end
                end
            end
        end
    endgenerate

    // B数据路由逻辑
    generate
        for (i_b = 0; i_b < 16; i_b = i_b + 1) begin : b_data_row
            for (j_b = 0; j_b < 16; j_b = j_b + 1) begin : b_data_col
                always @(*) begin
                    // 处理B数据
                    if (i_b == 0) begin
                        if (matrix_type == `m32n8k16) begin
                            if (j_b < 8) begin
                                b[i_b][j_b] = b_in[j_b];
                            end else if (j_b < 16) begin
                                b[i_b][j_b] = pe_b_out[15][j_b - 8];
                            end
                        end else begin
                            if (j_b < 16) begin
                                b[i_b][j_b] = b_in[j_b];
                            end
                        end
                    end else if (i_b == 8) begin
                        if (matrix_type == `m8n32k16) begin
                            if (j_b < 16) begin
                                b[i_b][j_b] = b_in[j_b + 16];
                            end
                        end else begin
                            b[i_b][j_b] = pe_b_out[i_b - 1][j_b];
                        end
                    end else begin
                        b[i_b][j_b] = pe_b_out[i_b - 1][j_b];
                    end
                end
            end
        end
    endgenerate

    /* // C数据路由逻辑
    generate
        for (i_c = 0; i_c < 16; i_c = i_c + 1) begin : c_data_row
            for (j_c = 0; j_c < 16; j_c = j_c + 1) begin : c_data_col
                always @(*) begin
                    // 处理C数据
                    case (matrix_type)
                        `m32n8k16: begin
                            if (i_c < 16 && j_c < 8) begin
                                c[i_c][j_c] = c_in[i_c][j_c];
                            end else if (i_c < 16 && j_c >= 8 && j_c < 16) begin
                                c[i_c][j_c] = c_in[i_c + 16][j_c - 8];
                            end
                            else begin
                                c[i_c][j_c] = 0;
                            end
                        end
                        `m16n16k16: begin
                            if (i_c < 16 && j_c < 16) begin
                                c[i_c][j_c] = c_in[i_c][j_c];
                            end
                            else begin
                                c[i_c][j_c] = 0;
                            end
                        end
                        `m8n32k16: begin
                            if (i_c < 8 && j_c < 16) begin
                                c[i_c][j_c] = c_in[i_c][j_c];
                            end else if (i_c >= 8 && i_c < 16 && j_c < 16) begin
                                c[i_c][j_c] = c_in[i_c - 8][j_c + 16];
                            end
                            else begin
                                c[i_c][j_c] = 0;
                            end
                        end
                    endcase
                end
            end
        end
    endgenerate */

/*     // C数据使能信号时序逻辑
    generate
        for (i_c = 0; i_c < 16; i_c = i_c + 1) begin : c_add_en_row
            for (j_c = 0; j_c < 16; j_c = j_c + 1) begin : c_add_en_col
                always @(posedge clk or posedge reset) begin
                    if (reset) begin
                        c_add_en[i_c][j_c] <= 0;
                    end else begin
                        case (matrix_type)
                            `m32n8k16: begin
                                if (i_c < 16 && j_c < 8) begin
                                    c_add_en[i_c][j_c] <= ~c_fifo_empty[i_c][j_c];
                                end else if (i_c < 16 && j_c >= 8 && j_c < 16) begin
                                    c_add_en[i_c][j_c] <= ~c_fifo_empty[i_c + 16][j_c - 8];
                                end else begin
                                    c_add_en[i_c][j_c] <= 0;
                                end
                            end
                            `m16n16k16: begin
                                if (i_c < 16 && j_c < 16) begin
                                    c_add_en[i_c][j_c] <= ~c_fifo_empty[i_c][j_c];
                                end else begin
                                    c_add_en[i_c][j_c] <= 0;
                                end
                            end
                            `m8n32k16: begin
                                if (i_c < 8 && j_c < 16) begin
                                    c_add_en[i_c][j_c] <= ~c_fifo_empty[i_c][j_c];
                                end else if (i_c >= 8 && i_c < 16 && j_c < 16) begin
                                    c_add_en[i_c][j_c] <= ~c_fifo_empty[i_c - 8][j_c + 16];
                                end else begin
                                    c_add_en[i_c][j_c] <= 0;
                                end
                            end
                        endcase
                    end
                end
            end
        end 
    endgenerate */

endmodule